// Copyright (C) 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef __GAPIL_RUNTIME_MAP_INC__
#define __GAPIL_RUNTIME_MAP_INC__

#include "map.h"
#include "hash.h"

#include "core/cc/assert.h"
#include "core/memory/arena/cc/arena.h"

#include <unordered_map>
#include <type_traits>

static_assert((GAPIL_MIN_MAP_SIZE & (GAPIL_MIN_MAP_SIZE - 1)) == 0, "Map size must be a power of 2");
static_assert((GAPIL_MAP_GROW_MULTIPLIER & (GAPIL_MAP_GROW_MULTIPLIER - 1)) == 0, "Map size must be a power of 2");

namespace gapil {

template<typename K, typename V, bool DENSE>
Map<K, V, DENSE>::Map(core::Arena* a) {
    GAPID_ASSERT_MSG(a != nullptr, "arena must not be null");
    auto buf = a->allocate(sizeof(Map<K, V, DENSE>::SVOAllocation), alignof(Map<K, V, DENSE>::SVOAllocation));
    memset(buf, 0, sizeof(Map<K, V, DENSE>::SVOAllocation));
    new(buf) Map<K, V, DENSE>::Allocation();
    Map<K, V, DENSE>::SVOAllocation* svo = static_cast<Map<K, V, DENSE>::SVOAllocation*>(buf);
    ptr = static_cast<Map<K, V, DENSE>::Allocation*>(buf);
    ptr->arena = reinterpret_cast<arena_t*>(a);
    ptr->capacity = GAPIL_MIN_MAP_SIZE;
    ptr->count = 0;
    ptr->elements = &svo->els;
    ptr->ref_count = 1;
}

template<typename K, typename V, bool DENSE>
Map<K, V, DENSE>::Map(const Map<K, V, DENSE>& s) {
    ptr = s.ptr;
    ptr->ref_count++;
}

template<typename K, typename V, bool DENSE>
Map<K, V, DENSE>::Map(Map<K, V, DENSE>&& s) {
    ptr = s.ptr;
    s.ptr = nullptr;
}

template<typename K, typename V, bool DENSE>
Map<K, V, DENSE>::~Map() {
    if (ptr != nullptr) { // note: nullptr is only valid in the case of a move
        ptr->release();
    }
}

template<typename K, typename V, bool DENSE>
Map<K, V, DENSE> Map<K, V, DENSE>::clone() const {
    auto out = Map(arena());
    for (auto it : *this) {
        out[it.first] = it.second;
    }
    return out;
}

template<typename K, typename V, bool DENSE>
Map<K, V, DENSE>& Map<K, V, DENSE>::operator = (const Map<K, V, DENSE>& other) {
    GAPID_ASSERT_MSG(other.ptr->ref_count > 0, "attempting to reference freed map");
    if (ptr != other.ptr) {
        ptr->release();
        ptr = other.ptr;
        ptr->reference();
    }
    return *this;
}

template<typename K, typename V, bool DENSE>
inline void Map<K, V, DENSE>::Allocation::reference() {
    GAPID_ASSERT_MSG(ref_count > 0, "Attempting to reference deleted map");
    ref_count++;
}

template<typename K, typename V, bool DENSE>
inline void Map<K, V, DENSE>::Allocation::release() {
    GAPID_ASSERT_MSG(ref_count > 0, "Attempting to release deleted map");
    ref_count--;
    if (ref_count > 0) {
        return;
    }
    auto arena = reinterpret_cast<core::Arena*>(this->arena);
    if (elements != nullptr) {
        clear_for_delete();
    }
    arena->free(this);
}

template<typename K, typename V, bool DENSE>
bool Map<K, V, DENSE>::Allocation::contains(K key) {
    return index(key, false) != nullptr;
}

template<typename T, bool DENSE>
typename std::enable_if<std::is_integral<T>::value, uint64_t>::type dense_hash(const T& t) {
    return static_cast<uint64_t>(t);
}


template<typename T, bool DENSE>
typename std::enable_if<!std::is_integral<T>::value, uint64_t>::type dense_hash(const T& t) {
    return 0;
}

const uint64_t k_valid_hash_bits = 0xFFFFFFFFFFF;

// Secondary hashes are generated by an LCG.
// https://en.wikipedia.org/wiki/Linear_congruential_generator
// The period of this LCG is k_valid_hash_bits. Meaning it will
// hit every possible hash before cycling. This is an important
// consideration.
// This is true because c=1.
// a = 69069 ( a-1 is divisible by 4)
// a = 69069 was chosen because (k_valid_hash_bits * a) will not
// overflow a 64-bit number. If you want to change this, make sure
// you understand these implications.
inline uint64_t next_hash(uint64_t hash) {
    return ((hash * 69069) + 1) & k_valid_hash_bits;
}


template<typename K, typename V, bool DENSE>
V* Map<K, V, DENSE>::Allocation::index(K key, bool insert) {
    // TODO(awoloszyn): Template tricks to move this to compile-time
    if (DENSE) {
        uint64_t hash = dense_hash<K, DENSE>(key);
        if (hash >= capacity) {
            if (insert) {
                if (elements == nullptr) {
                    // Start the dense map larger
                    capacity = GAPIL_MIN_MAP_SIZE * 16;
                    // Capacity may be less than hash, but it still must be a power of 2
                    while(capacity < hash) {
                        capacity >>= 1;
                    }
                    elements = gapil_alloc(arena, sizeof(element) * capacity, alignof(V));
                    for (uint64_t i = 0; i < capacity; ++i) {
                        this->els()[i].used = GAPIL_MAP_ELEMENT_EMPTY;
                    }
                } else {
                    auto oldElements = this->els();
                    auto oldCapacity = capacity;
                    capacity = capacity * GAPIL_MAP_GROW_MULTIPLIER * 2;
                    // Capacity may be less than hash, but it still must be a power of 2
                    // Grow the dense map faster, we probably don't have to create quite as many
                    while(capacity < hash) {
                        capacity >>= 1;
                    }
                    elements = gapil_alloc(arena, sizeof(element) * capacity, alignof(V));
                    for (uint64_t i = 0; i < capacity; ++i) {
                        els()[i].used = GAPIL_MAP_ELEMENT_EMPTY;
                    }
                    auto new_elements = els();
                    for (uint64_t i = 0; i < oldCapacity; ++i) {
                        if (oldElements[i].used == GAPIL_MAP_ELEMENT_FULL) {
                            new(&new_elements[i].second) V(std::move(oldElements[i].second));
                            new(&new_elements[i].first) K(std::move(oldElements[i].first));
                            new_elements[i].used = GAPIL_MAP_ELEMENT_FULL;
                            oldElements[i].second.~V();
                            oldElements[i].first.~K();
                        }
                    }
                    if (oldCapacity != GAPIL_MIN_MAP_SIZE) {
                        gapil_free(arena, oldElements);
                    }
                }
            } else {
                return nullptr;
            }
        }
        auto elems = els();
        if (elems[hash].used == GAPIL_MAP_ELEMENT_FULL) {
            return &elems[hash].second;
        } else if (insert) {
            count++;
            auto a = reinterpret_cast<core::Arena*>(arena);
            inplace_new(&elems[hash].second, a);
            inplace_new(&elems[hash].first, a, key);
            elems[hash].used = GAPIL_MAP_ELEMENT_FULL;
            return &elems[hash].second;
        }
        return nullptr;
    }
    auto hasher = gapil::hash<K>{};
    auto eq = std::equal_to<K>{};
    uint64_t hash = hasher(key) & k_valid_hash_bits;

    auto elems = els();

    for (uint64_t i = 0; i < capacity; ++i, hash = next_hash(hash)) {
        bool leave = false;
        uint64_t lookup_pos = hash & (capacity - 1);
        switch(elems[lookup_pos].used) {
            case GAPIL_MAP_ELEMENT_EMPTY:
                leave = true;
                break;
            case GAPIL_MAP_ELEMENT_USED:
                continue;
            case GAPIL_MAP_ELEMENT_FULL:
                if (eq(key, elems[lookup_pos].first)) {
                    return &elems[lookup_pos].second;
                }
        }
        if (leave) {
            break;
        }
    }

    // storageBucket assumes there is at least one open cell.
    // Make sure before you run this, that is the case.
    auto storageBucket = [&](uint64_t h) {
        h = h & k_valid_hash_bits;
        auto elems = els();
        for (uint64_t i = 0; i < capacity; ++i, h = next_hash(h)) {
            uint64_t x = h & (capacity - 1);
            if (elems[x].used != GAPIL_MAP_ELEMENT_FULL) {
                return x;
            }
        }
        return uint64_t(0);
    };

    if (insert) {
        bool resize = (elements == nullptr);
        resize = resize || ((float)count / (float)capacity) > GAPIL_MAP_MAX_CAPACITY;

        if (resize) {
            if (elements == nullptr) {
                capacity = GAPIL_MIN_MAP_SIZE;
                elements = gapil_alloc(arena, sizeof(element) * GAPIL_MIN_MAP_SIZE, alignof(V));
                for (uint64_t i = 0; i < capacity; ++i) {
                    this->els()[i].used = GAPIL_MAP_ELEMENT_EMPTY;
                 }
            } else {
                 auto oldElements = this->els();
                 auto oldCapacity = capacity;

                 capacity = capacity * GAPIL_MAP_GROW_MULTIPLIER;
                 elements = gapil_alloc(arena, sizeof(element) * capacity, alignof(V));
                 for (uint64_t i = 0; i < capacity; ++i) {
                    els()[i].used = GAPIL_MAP_ELEMENT_EMPTY;
                 }
                 auto new_elements = els();
                 for (uint64_t i = 0; i < oldCapacity; ++i) {
                     if (oldElements[i].used == GAPIL_MAP_ELEMENT_FULL) {
                        uint64_t bucket_location = storageBucket(hasher(oldElements[i].first));
                        new(&new_elements[bucket_location].second) V(std::move(oldElements[i].second));
                        new(&new_elements[bucket_location].first) K(std::move(oldElements[i].first));
                        new_elements[bucket_location].used = GAPIL_MAP_ELEMENT_FULL;
                        oldElements[i].second.~V();
                        oldElements[i].first.~K();
                     }
                 }
                 if (oldCapacity != GAPIL_MIN_MAP_SIZE) {
                    gapil_free(arena, oldElements);
                 }
             }
        }

        uint64_t bucket_location = storageBucket(hasher(key));
        auto a = reinterpret_cast<core::Arena*>(arena);
        inplace_new(&els()[bucket_location].second, a);
        inplace_new(&els()[bucket_location].first, a, key);
        els()[bucket_location].used = GAPIL_MAP_ELEMENT_FULL;
        count++;

        return &els()[bucket_location].second;
    }

    return nullptr;
}

template<typename K, typename V, bool DENSE>
V Map<K, V, DENSE>::Allocation::lookup(K key) {
    V* v = index(key, false);
    return *v;
}

template<typename K, typename V, bool DENSE>
void Map<K, V, DENSE>::Allocation::remove(K key) {
    // TODO(awoloszyn): Template tricks to move this to compile-time
    if (DENSE) {
        uint64_t hash = dense_hash<K, DENSE>(key);
        auto elems = els();
        if (hash < capacity &&
                elems[hash].used == GAPIL_MAP_ELEMENT_FULL) {
            elems[hash].used = GAPIL_MAP_ELEMENT_EMPTY;
            elems[hash].first.~K();
            elems[hash].second.~V();
            --count;
        }
        return;
    }
    auto hasher = gapil::hash<K>{};
    auto eq = std::equal_to<K>{};
    uint64_t hash = hasher(key) & k_valid_hash_bits;
    auto elems = els();

    for (uint64_t i = 0; i < capacity; ++i, hash = next_hash(hash)) {
        uint64_t lookup_pos = hash & (capacity - 1);
        switch(elems[lookup_pos].used) {
            case GAPIL_MAP_ELEMENT_EMPTY:
                return;
            case GAPIL_MAP_ELEMENT_USED:
                continue;
            case GAPIL_MAP_ELEMENT_FULL:
                if (eq(key, elems[lookup_pos].first)) {
                    elems[lookup_pos].used = GAPIL_MAP_ELEMENT_USED;
                    elems[lookup_pos].first.~K();
                    elems[lookup_pos].second.~V();
                    --count;
                    return;
                }
        }
    }
}

template<typename K, typename V, bool DENSE>
void Map<K, V, DENSE>::Allocation::clear_for_delete() {
    auto elems = els();
    for (uint64_t i = 0; i < capacity && count > 0; ++i) {
        switch(elems[i].used) {
            case GAPIL_MAP_ELEMENT_EMPTY:
            case GAPIL_MAP_ELEMENT_USED:
                continue;
            case GAPIL_MAP_ELEMENT_FULL:
                elems[i].first.~K();
                elems[i].second.~V();
                --count;
        }
    }
    if (capacity != GAPIL_MIN_MAP_SIZE) {
        gapil_free(arena, elements);
        capacity = 0;
        elements = nullptr;
    }
    count = 0;
}

// The warning "Wclass-memaccess" added with GCC 8+ is raised by memset.
#if defined __GNUC__ && __GNUC__ >= 8
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wclass-memaccess"
#endif

template<typename K, typename V, bool DENSE>
void Map<K, V, DENSE>::Allocation::clear_keep() {
    auto elems = els();
    for (uint64_t i = 0; i < capacity && count > 0; ++i) {
        switch(elems[i].used) {
            case GAPIL_MAP_ELEMENT_EMPTY:
            case GAPIL_MAP_ELEMENT_USED:
                continue;
            case GAPIL_MAP_ELEMENT_FULL:
                elems[i].first.~K();
                elems[i].second.~V();
                --count;
        }
    }

    memset(elems, 0x00, sizeof(*elems) * capacity);
    count = 0;
}

}  // namespace gapil

#if defined __GNUC__ && __GNUC__ >= 8
#pragma GCC diagnostic pop
#endif

#endif  // __GAPIL_RUNTIME_MAP_INC__
